{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='1'></a>\n",
    "# Import modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import keras.backend as K"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='4'></a>\n",
    "# Model Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "K.set_learning_phase(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input/Output resolution\n",
    "RESOLUTION = 64 # 64x64, 128x128, 256x256\n",
    "assert (RESOLUTION % 64) == 0, \"RESOLUTION should be 64, 128, 256\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Architecture configuration\n",
    "arch_config = {}\n",
    "arch_config['IMAGE_SHAPE'] = (RESOLUTION, RESOLUTION, 3)\n",
    "arch_config['use_self_attn'] = True\n",
    "arch_config['norm'] = \"instancenorm\" # instancenorm, batchnorm, layernorm, groupnorm, none\n",
    "arch_config['model_capacity'] = \"standard\" # standard, lite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='5'></a>\n",
    "# Define models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from networks.faceswap_gan_model import FaceswapGANModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model = FaceswapGANModel(**arch_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='6'></a>\n",
    "# Load Model Weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model weights files are successfully loaded\n"
     ]
    }
   ],
   "source": [
    "model.load_weights(path=\"./models\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='12'></a>\n",
    "# Video Conversion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from converter.video_converter import VideoConverter\n",
    "from detector.face_detector import MTCNNFaceDetector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mtcnn_weights_dir = \"./mtcnn_weights/\"\n",
    "\n",
    "fd = MTCNNFaceDetector(sess=K.get_session(), model_path=mtcnn_weights_dir)\n",
    "vc = VideoConverter()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "vc.set_face_detector(fd)\n",
    "vc.set_gan_model(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Video conversion configuration\n",
    "\n",
    "\n",
    "- `use_smoothed_bbox`: \n",
    "    - Boolean. Whether to enable smoothed bbox.\n",
    "- `use_kalman_filter`: \n",
    "    - Boolean. Whether to enable Kalman filter.\n",
    "- `use_auto_downscaling`:\n",
    "    - Boolean. Whether to enable auto-downscaling in face detection (to prevent OOM error).\n",
    "- `bbox_moving_avg_coef`: \n",
    "    - Float point between 0 and 1. Smoothing coef. used when use_kalman_filter is set False.\n",
    "- `min_face_area`:\n",
    "    - int x int. Minimum size of face. Detected faces smaller than min_face_area will not be transformed.\n",
    "- `IMAGE_SHAPE`:\n",
    "    - Input/Output resolution of the GAN model\n",
    "- `kf_noise_coef`:\n",
    "    - Float point. Increase by 10x if tracking is slow. Decrease by 1/10x if trakcing works fine but jitter occurs.\n",
    "- `use_color_correction`: \n",
    "    - String of \"adain\", \"adain_xyz\", \"hist_match\", or \"none\". The color correction method to be applied.\n",
    "- `detec_threshold`: \n",
    "    - Float point between 0 and 1. Decrease its value if faces are missed. Increase its value to reduce false positives.\n",
    "- `roi_coverage`: \n",
    "    - Float point between 0 and 1 (exclusive). Center area of input images to be cropped (Suggested range: 0.85 ~ 0.95)\n",
    "- `enhance`: \n",
    "    - Float point. A coef. for contrast enhancement in the region of alpha mask (Suggested range: 0. ~ 0.4)\n",
    "- `output_type`: \n",
    "    - Layout format of output video: 1. [ result ], 2. [ source | result ], 3. [ source | result | mask ]\n",
    "- `direction`: \n",
    "    - String of \"AtoB\" or \"BtoA\". Direction of face transformation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "options = {\n",
    "    # ===== Fixed =====\n",
    "    \"use_smoothed_bbox\": True,\n",
    "    \"use_kalman_filter\": True,\n",
    "    \"use_auto_downscaling\": False,\n",
    "    \"bbox_moving_avg_coef\": 0.65,\n",
    "    \"min_face_area\": 35 * 35,\n",
    "    \"IMAGE_SHAPE\": model.IMAGE_SHAPE,\n",
    "    # ===== Tunable =====\n",
    "    \"kf_noise_coef\": 3e-3,\n",
    "    \"use_color_correction\": \"hist_match\",\n",
    "    \"detec_threshold\": 0.7,\n",
    "    \"roi_coverage\": 0.9,\n",
    "    \"enhance\": 0.,\n",
    "    \"output_type\": 3,\n",
    "    \"direction\": \"AtoB\",\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Start video conversion\n",
    "\n",
    "\n",
    "- `input_fn`: \n",
    "    - String. Input video path.\n",
    "- `output_fn`: \n",
    "    - String. Output video path.\n",
    "- `duration`: \n",
    "    - None or a non-negative float tuple: (start_sec, end_sec). Duration of input video to be converted\n",
    "    - e.g., setting `duration = (5, 7.5)` outputs a 2.5-sec-long video clip corresponding to 5s ~ 7.5s of the input video."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_fn = \"INPUT_VIDEO.mp4\"\n",
    "output_fn = \"OUTPUT_VIDEO.mp4\"\n",
    "duration = None "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[MoviePy] >>>> Building video OUTPUT_VIDEO.mp4\n",
      "[MoviePy] Writing video OUTPUT_VIDEO.mp4\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████| 20/20 [00:12<00:00,  1.48it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[MoviePy] Done.\n",
      "[MoviePy] >>>> Video ready: OUTPUT_VIDEO.mp4 \n",
      "\n"
     ]
    }
   ],
   "source": [
    "vc.convert(input_fn=input_fn, output_fn=output_fn, options=options, duration=duration)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
